#!/usr/bin/env python3
"""
AES-256密钥生成工具 - V4.0零信任架构
用于生成模型加密的强密钥

使用方法:
python generate_key.py

生成的密钥将用于:
1. 加密模型文件
2. 作为环境变量 MODEL_ENCRYPTION_KEY 注入到生产环境
3. 解密时在内存中验证密钥有效性

⚠️ 安全警告:
- 生成的密钥必须安全存储
- 建议使用专业的密钥管理系统(KMS)
- 生产环境中禁止硬编码密钥
"""

import os
import sys
import secrets
import base64
import hashlib
import getpass
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


def generate_secure_key():
    """生成加密安全的256位AES密钥"""
    # 使用cryptography.fernet生成32字节(256位)密钥
    key = Fernet.generate_key()
    return key


def derive_key_from_password(password: str, salt: bytes = None) -> bytes:
    """
    从密码派生密钥 (PBKDF2)

    Args:
        password: 用户密码
        salt: 盐值，如果为None则随机生成

    Returns:
        32字节的派生密钥
    """
    if salt is None:
        # 生成随机盐值
        salt = os.urandom(16)

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,  # 高迭代次数增强安全性
    )

    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
    return key, salt


def validate_key_strength(key: bytes) -> bool:
    """验证密钥强度"""
    # Fernet密钥固定为32字节(base64编码后44字符)
    if len(key) != 44:
        return False

    try:
        # 尝试创建Fernet实例验证密钥格式
        Fernet(key)
        return True
    except Exception:
        return False


def display_security_warning():
    """显示安全警告"""
    print("=" * 70)
    print("🔐 AES-256 密钥生成工具 - V4.0零信任架构")
    print("=" * 70)
    print("⚠️  重要安全警告:")
    print("   1. 请将生成的密钥安全保存到密钥管理系统")
    print("   2. 不要将密钥提交到版本控制系统")
    print("   3. 生产环境建议使用专业的KMS服务")
    print("   4. 密钥泄露将导致模型知识产权完全暴露")
    print("=" * 70)


def generate_key_with_entropy_check():
    """生成具有高熵值的密钥"""
    print("正在生成高熵值AES-256密钥...")

    # 收集系统熵源
    entropy_sources = [
        secrets.token_bytes(32),  # 加密安全随机字节
        str(os.getpid()).encode(),  # 进程ID
        str(id(object())).encode(),  # 对象ID
    ]

    # 添加用户交互熵（可选）
    try:
        print("请输入一些随机字符增强熵值 (按Enter跳过):")
        user_input = input().strip().encode()
        if user_input:
            entropy_sources.append(user_input)
    except KeyboardInterrupt:
        print("\n跳过用户熵输入")

    # 组合熵源并生成最终密钥
    combined_entropy = b''.join(entropy_sources)
    final_key = generate_secure_key()

    # 验证密钥强度
    if not validate_key_strength(final_key):
        raise RuntimeError("生成的密钥强度验证失败")

    return final_key


def save_key_to_file(key: bytes, filename: str = "model_encryption.key"):
    """安全保存密钥到文件"""
    try:
        with open(filename, 'wb') as f:
            f.write(key)

        # 设置文件权限为仅所有者可读写
        os.chmod(filename, 0o600)
        print(f"✅ 密钥已保存到: {filename}")
        print(f"   文件权限已设置为: 仅所有者可访问")

    except Exception as e:
        print(f"❌ 保存密钥失败: {e}")
        return False

    return True


def display_key_info(key: bytes):
    """显示密钥信息"""
    key_str = key.decode('utf-8')

    print("\n🔑 生成的密钥信息:")
    print(f"   密钥长度: {len(key)} 字节 ({len(key) * 8} 位)")
    print(f"   算法: AES-256-CBC (Fernet)")
    print(f"   Base64编码长度: {len(key_str)} 字符")

    # 计算密钥指纹
    fingerprint = hashlib.sha256(key).hexdigest()[:16]
    print(f"   密钥指纹: {fingerprint}")

    print("\n📋 环境变量配置:")
    print(f"   MODEL_ENCRYPTION_KEY={key_str}")

    print("\n⚠️  重要提醒:")
    print("   1. 立即将上述环境变量添加到生产环境")
    print("   2. 或者使用专业的密钥管理系统")
    print("   3. 不要在代码中硬编码此密钥")
    print("   4. 定期轮换密钥（建议每6个月）")


def test_key_functionality(key: bytes):
    """测试密钥功能"""
    print("\n🧪 测试密钥功能...")

    try:
        fernet = Fernet(key)

        # 测试加密/解密
        test_data = b"EndoSight-UC V4.0 Security Test"
        encrypted = fernet.encrypt(test_data)
        decrypted = fernet.decrypt(encrypted)

        if decrypted == test_data:
            print("✅ 密钥功能测试通过")
            return True
        else:
            print("❌ 密钥功能测试失败")
            return False

    except Exception as e:
        print(f"❌ 密钥测试异常: {e}")
        return False


def main():
    """主函数"""
    display_security_warning()

    try:
        # 生成密钥
        key = generate_key_with_entropy_check()

        # 显示密钥信息
        display_key_info(key)

        # 测试密钥功能
        if not test_key_functionality(key):
            print("❌ 密钥测试失败，请重新生成")
            sys.exit(1)

        # 询问是否保存到文件
        try:
            save_to_file = input("\n是否将密钥保存到文件? (y/n): ").strip().lower()
            if save_to_file in ['y', 'yes', '是']:
                filename = input("请输入文件名 (默认: model_encryption.key): ").strip()
                if not filename:
                    filename = "model_encryption.key"
                save_key_to_file(key, filename)
        except KeyboardInterrupt:
            print("\n跳过文件保存")

        print("\n🎉 密钥生成完成!")
        print("📝 下一步:")
        print("   1. 运行: python encrypt_model.py")
        print("   2. 配置生产环境变量 MODEL_ENCRYPTION_KEY")
        print("   3. 启动V4.0安全推理服务")

    except KeyboardInterrupt:
        print("\n\n⚠️  密钥生成已取消")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 密钥生成失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()